import numpy as np
from enum import Enum
from ._base import EquatorialWave
from ..physics.constant import PI, EARTH_GRAVITY, ROSSBY_PARAMETER_ON_EQUATOR
from pint import Quantity
from .. import unit
from ..physics.unit import _handle_quantity




def dimensionless_distribution(x: np.ndarray, y: np.ndarray, t: float,
                               k: float, omega: float, n:int=1
                               ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """计算赤道波动的 u, v, phi 空间分布。需要传入时空坐标与波数、频率等参数。
    波数与频率需要先通过频散关系得到。
    返回三个数组，分别对应 u, v, phi 的空间分布。

    Calculate the spatial distribution of u, v, and phi for equatorial waves.
    The wave number and frequency should be obtained from the dispersion relation.
    Returns three arrays corresponding to the spatial distributions of u, v,
    and phi.

    Parameters
    ----------
    x, y : np.ndarray
        空间的无量纲坐标，可以都是一维或二维的数组。

        The dimensionless spatial coordinates, which can be either 1D or 2D
        arrays.
    t : float
        时间  Time
    k : float
        波数  Wave number
    omega : float
        圆频率  Angular frequency
    n : int, optional
        Hermite多项式的阶数

        The order of the Hermite polynomial, by default 1

    Returns
    -------
    tuple[np.ndarray, np.ndarray, np.ndarray]
        u, v, phi 的空间分布

        The spatial distributions of u, v, and phi.

    """
    if x.ndim == 1 and y.ndim == 1:
        x, y = np.meshgrid(x, y)
    elif x.ndim == 2 and y.ndim == 2:
        pass
    else:
        raise ValueError("x and y must both be 1D or 2D arrays.")

    from scipy.special import hermite
    ey = np.exp(-y**2/2)
    if n >= 1:
        v = -(omega**2-k**2) * ey * hermite(n, True)(y) * np.sin(k*x + omega*t)
        u = ey * (0.5 * (omega - k) * hermite(n+1, True)(y) + n*(omega + k) *\
                   hermite(n-1, True)(y)) * np.cos(k*x + omega*t)
        phi = ey * (0.5 * (omega - k) * hermite(n+1, True)(y) - n*(omega + k) *\
                     hermite(n-1, True)(y)) * np.cos(k*x + omega*t)
    elif n == 0:
        v = -2 * (omega+k) * ey * np.sin(k*x + omega*t)
        u = 2 * y * ey * np.cos(k*x + omega*t)
        phi = 2 * y * ey * np.cos(k*x + omega*t)
    elif n == -1:
        v = np.zeros_like(x)
        u = ey * np.cos(k*x + omega*t)
        phi = ey * np.cos(k*x + omega*t)
    else:
        raise ValueError("n must be >= -1")

    return u, v, phi


